import logging
from logging.handlers import TimedRotatingFileHandler

LOG_FORMAT = "%(asctime)s [%(levelname)s] %(message)s"


class ConsoleHandler(logging.StreamHandler):
    GRAY8 = "38;5;8"
    GRAY7 = "38;5;7"
    ORANGE = "33"
    RED = "31"
    WHITE = "97"
    PURPLE = "35"
    BLUE = "34"
    GREEN = "32"
    YELLOW = "33"
    CYAN = "36"

    # 新增 __init__ 方法，支持传入 formatter
    def __init__(self, stream=None, formatter=logging.Formatter(LOG_FORMAT)):
        super().__init__(stream)
        if formatter is not None:
            self.setFormatter(formatter)
        # 设置默认级别INFO
        self.setLevel(logging.INFO)
        # 设置默认日志颜色格式映射
        self.level_color_map = {
            logging.DEBUG: self.BLUE,
            logging.INFO: self.WHITE,
            logging.WARNING: self.ORANGE,
            logging.ERROR: self.RED,
            logging.CRITICAL: self.PURPLE
        }

    def emit(self, record):
        try:
            msg = self.format(record)
            csi = f"{chr(27)}["  # control sequence introducer
            color = self.level_color_map.get(record.levelno, self.WHITE)

            self.stream.write(f"{csi}{color}m{msg}{csi}m\n")
            self.flush()
        except RecursionError:
            raise
        except Exception:
            self.handleError(record)


class FileHandler(TimedRotatingFileHandler):
    def __init__(self, filename, when='midnight', interval=1, backupCount=7,
                 encoding=None, delay=False, formatter=logging.Formatter(LOG_FORMAT)):
        super().__init__(filename, when=when, interval=interval,
                         backupCount=backupCount, encoding=encoding, delay=delay)

        if formatter is not None:
            self.setFormatter(formatter)

        # 设置默认级别DEBUG
        self.setLevel(logging.DEBUG)

    def emit(self, record):
        try:
            # 只格式化消息，不加颜色
            msg = self.format(record)
            self.stream.write(f"{msg}\n")
            self.flush()
        except RecursionError:
            raise
        except Exception:
            self.handleError(record)
